package com.yuan.aiagent.AIInvoke.rag;

import jakarta.annotation.Resource;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.rag.Query;
import org.springframework.ai.rag.preretrieval.query.expansion.MultiQueryExpander;
import org.springframework.context.annotation.Configuration;
import org.springframework.stereotype.Component;
import org.stringtemplate.v4.ST;

import java.util.List;

@Component
public class MultiQueryExpanderDemo {

    private ChatClient.Builder chatClientBuilder;

    public MultiQueryExpanderDemo(ChatModel ollamaChatModel){
        this.chatClientBuilder = ChatClient.builder(ollamaChatModel);
    }

    public List<Query> expandQuery(String query){
        MultiQueryExpander queryExpander = MultiQueryExpander.builder()
                .chatClientBuilder(chatClientBuilder)
                .numberOfQueries(3)
                .build();
        List<Query> queries = queryExpander.expand(new Query(query));
        return queries;
    }

}
